Window Functions

Unlike traditional groupBy() aggregations that collapse multiple rows into a single summary row, Window Functions compute values over a group of rows (a "window") while preserving the identity of each individual row in the final output.

graph TD
    subgraph Dataset["DataFrame Row Stream"]
        direction TB
        R1["Row A (Dept: Sales, Salary: 5000)"]
        R2["Row B (Dept: Sales, Salary: 6000)"]
        R3["Row C (Dept: Eng,   Salary: 8000)"]
    end
    subgraph Windows["Window Partitions (partitionBy)"]
        direction TB
        subgraph Partition1["Sales Department (orderBy salary desc)"]
            P1_R1["1. Bob (6000)"]
            P1_R2["2. Alice (5000)"]
        end
        subgraph Partition2["Engineering Department (orderBy salary desc)"]
            P2_R1["1. Eva (9500)"]
            P2_R2["2. David (8000)"]
        end
    end
    Dataset --> Windows
    style Partition1 fill:#eff6ff,stroke:#2563eb,stroke-width:1px;
    style Partition2 fill:#faf5ff,stroke:#9333ea,stroke-width:1px;

Defining a Window Specification

To write a window function, you must first construct a Window Specification using pyspark.sql.expressions.Window:

from pyspark.sql.expressions import Window

windowSpec = Window \
    .partitionBy("department") \
    .orderBy(col("salary").desc())

• partitionBy("col"): Defines the grouping boundary (equivalent to GROUP BY but doesn't collapse rows).
• orderBy("col"): Defines how rows are sorted inside each group partition.

Core Window Functions Reference

Import these functions from pyspark.sql.functions:

PySpark Code Example: Rankings & Running Sums

Here is a complete script demonstrating ranking employees within departments and calculating running totals:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.expressions import Window

# 1. Setup Spark
spark = SparkSession.builder \
    .appName("Window Functions") \
    .master("local[*]") \
    .getOrCreate()

# 2. Sample Employee Dataset
employee_data = [
    ("Sales", "Alice", 5000),
    ("Sales", "Bob", 6000),
    ("Sales", "Charlie", 5000),
    ("Engineering", "David", 8000),
    ("Engineering", "Eva", 9500),
    ("Engineering", "Frank", 8000)
]
columns = ["department", "name", "salary"]
df = spark.createDataFrame(employee_data, columns)

# 3. Create Window Specifications
# Rank Window
rank_window = Window.partitionBy("department").orderBy(F.col("salary").desc())

# Running Sum Window (Includes rows from start of partition to current row)
running_sum_window = Window.partitionBy("department") \
    .orderBy("salary") \
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 4. Calculate Rankings
# Observe the difference between row number, rank, and dense rank on matching salaries (5000 and 8000)!
ranked_df = df.withColumn("row_num", F.row_number().over(rank_window)) \
              .withColumn("rank", F.rank().over(rank_window)) \
              .withColumn("dense_rank", F.dense_rank().over(rank_window))

ranked_df.show()

# 5. Calculate Running Sum within Department
running_sum_df = df.withColumn("running_total", F.sum("salary").over(running_sum_window))
running_sum_df.show()